import os

import torch
import torchvision
import torchvision.transforms as transforms
import numpy
from PIL import Image


class ANN:
    def __init__(self):
        self.a0 = torch.rand((1, 784))
        self.w0 = torch.rand((784, 1000)) * 0.001
        self.b0 = torch.rand((1, 1000)) * 0.001
        self.x0 = torch.rand((1, 1000))
        self.a1 = torch.rand((1, 1000))
        self.w1 = torch.rand((1000, 10)) * 0.001
        self.b1 = torch.rand((1, 10)) * 0.001
        self.x1 = torch.rand((1, 10))
        self.y = torch.rand((1, 10))

        self.dy = torch.rand((1, 10))
        self.dx1 = torch.rand((1, 10))
        self.dw1 = torch.rand((1000, 10))
        self.db1 = torch.rand((1, 10))
        self.da1 = torch.rand((1, 1000))
        self.dx0 = torch.rand((1, 1000))
        self.dw0 = torch.rand((784, 1000))
        self.db0 = torch.rand((1, 1000))

        self.alpha = 0.00001
    def load(self):
        file = open("a.txt")
        for i in range(784):
            line = file.readline().split()
            for j in range(1000):
                self.w0[i][j] = float(line[j])
        line = file.readline().split()
        for i in range(1000):
            self.b0[0][i] = float(line[i])
        for i in range(1000):
            line = file.readline().split()
            for j in range(10):
                self.w1[i][j] = float(line[j])
        line = file.readline().split()
        for i in range(10):
            self.b1[0][i] = float(line[i])

    def save(self):
        file = open("a.txt", "w")
        for i in range(784):
            for j in range(1000):
                file.write(str(float(self.w0[i][j])) + " ")
            file.write("\n")
        for i in range(1000):
            file.write(str(float(self.b0[0][i])) + " ")
        file.write("\n")
        for i in range(1000):
            for j in range(10):
                file.write(str(float(self.w1[i][j])) + " ")
            file.write("\n")
        for i in range(10):
            file.write(str(float(self.b1[0][i])) + " ")
        file.write("\n")

    def train(self, dataSet):
        n = 0
        L = 0
        for i, (image, label) in enumerate(dataSet):
            # self.a0 = image.reshape((1, 784))
            # self.x0 = torch.mm(self.a0, self.w0) + self.b0
            # self.a1 = torch.where(self.x0 > 0, self.x0, float(0))
            # self.x1 = torch.mm(self.a1, self.w1) + self.b1
            # self.y = torch.where(self.x1 > 0, self.x1, float(0))
            #
            # self.dy = self.y.clone()
            # self.dy[0][label] -= 1
            # self.dx1 = self.dy * torch.where(self.x1 > 0, float(1), float(0))
            # self.dw1 = self.dx1.repeat((1000,1)) * self.a1.T.repeat((1,10))
            # self.db1 = self.dx1.clone()
            # self.da1 = torch.mm(self.dx1, self.w1.T)
            # self.dx0 = self.da1 * torch.where(self.x0 > 0, float(1), float(0))
            # self.dw0 = self.dx0.repeat((784,1)) * self.a0.T.repeat((1,1000))
            # self.db0 = self.dx0.clone()

            self.a0 = image.reshape((1, 784))
            self.x0 = torch.mm(self.a0, self.w0) + self.b0
            self.a1 = torch.tanh(self.x0)
            self.x1 = torch.mm(self.a1, self.w1) + self.b1
            self.y = torch.tanh(self.x1)

            self.dy = self.y.clone()
            self.dy[0][label] -= 1
            self.dx1 = (1 - torch.tanh(self.x1) * torch.tanh(self.x1)) * self.dy
            self.dw1 = self.dx1.repeat((1000, 1)) * self.a1.T.repeat((1, 10))
            self.db1 = self.dx1.clone()
            self.da1 = torch.mm(self.dx1, self.w1.T)
            self.dx0 = (1 - torch.tanh(self.x0) * torch.tanh(self.x0)) * self.da1
            self.dw0 = self.dx0.repeat((784, 1)) * self.a0.T.repeat((1, 1000))
            self.db0 = self.dx0.clone()

            self.w0 -= self.dw0 * self.alpha
            self.b0 -= self.db0 * self.alpha
            self.w1 -= self.dw1 * self.alpha
            self.b1 -= self.db1 * self.alpha

            L += self.dy @ self.dy.T

            if self.y[0].argmax() == int(label[0]):
                n += 1

            if int(i) % 500 == 0:
                print(f"{n}/{int(i)+1}:{n/(int(i)+1)}")
                print(f"LOSS:{L/n}")
                print()



if __name__ == "__main__":

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    train_dataset = torchvision.datasets.MNIST(root="data", train=True, download=False, transform=transforms.ToTensor())
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=0)

    ann = ANN()
    ann.load()
    while True:
        train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=0)
        ann.train(train_loader)
        ann.save()
    # for i, (image, label) in enumerate(train_loader):
    #     out = snn.classify(tensor2time(image))
    #     count1[int(label)] += 1
    #     if out == int(label):
    #         # pil_image = transforms.ToPILImage()(image[0][0])
    #         # pil_image.save(f"wrong-train/{int(label)}-{out}-{count0[int(label)]}.png")
    #         count0[int(label)]+=1
    #     if int(i) % 10000 == 0 and int(i) != 0:
    #         print(int(i))
    #         print([count0[i]/count1[i] for i in range(10)])
    # while True:
    #     train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=0)
    #     snn.load()
    #     snn.train(train_loader)
    #     snn.save()
